#include <stdio.h>
#include <stdlib.h>
#include <cuda_runtime.h>

__global__ void global_scan(float *d_out, float * d_in){
    int idx = threadIdx.x;
    d_out[idx] = d_in[idx];
    __syncthreads();

    for(int stride = 1; stride < blockDim.x; stride <<= 1){
        if(idx - stride >= 0){
            float out = d_out[idx] + d_out[idx - stride];
            __syncthreads();
            d_out[idx] = out;
        }
        __syncthreads();
    }
}
// Hillis-steele算法
__global__ void sharemem_scan(float *d_out, float *d_in){
    extern __shared__ float sdata[];
    int idx = threadIdx.x;
    sdata[idx] = d_in[idx];
    __syncthreads();
    
    for(int stride = 1; stride < blockDim.x; stride <<= 1){
        if(idx >= stride){
            float temp = sdata[idx] + sdata[idx - stride];
            __syncthreads();
            sdata[idx] = temp;
        }
        __syncthreads();
    }
    d_out[idx] = sdata[idx];
}

__global__ void blelloch_scan(float *d_out, float *d_in) {
    extern __shared__ float temp[];
    int tid = threadIdx.x;
    float tail = 0.0f;
    
    temp[tid] = d_in[tid];
    __syncthreads();
    // up sweep
    for(int s = 1; s < blockDim.x; s <<= 1) {
        int index = tid * (s << 1);
        if(index + (s << 1) - 1 < blockDim.x) {
            temp[index + (s << 1) - 1] += temp[index + s - 1];
        }
        __syncthreads();
    }
    // reset the last element
    if(tid == blockDim.x - 1) {
        tail = temp[tid];
        temp[tid] = 0.0f;
        
    }
    __syncthreads();
    // down sweep
    for(int s = blockDim.x >> 1; s > 0; s >>= 1) {
        int index = tid * (s << 1);
        if(index + (s << 1) - 1 < blockDim.x) {
            float t = temp[index + s - 1];
            temp[index + s - 1] = temp[index + (s << 1) - 1];
            temp[index + (s << 1) - 1] += t;
        }
        __syncthreads();
    }
    
    if(tid < blockDim.x) {
        d_out[tid] = (tid < blockDim.x - 1) ? temp[tid + 1] : tail;
    }
}

int main(){
    const int ARRAY_SIZE = 32;
    const int ARRAY_BYTES = ARRAY_SIZE * sizeof(float);

    float h_in[ARRAY_SIZE];
	for(int i = 0; i < ARRAY_SIZE; i++){
		h_in[i] = (float)i;
	}
    float h_out[ARRAY_SIZE], h_out2[ARRAY_SIZE], h_out3[ARRAY_SIZE];

    float *d_in, *d_out, *d_in2, *d_out2, *d_in3, *d_out3;
	// allocate gpu memory
	cudaMalloc((float**)&d_in, ARRAY_BYTES);
	cudaMalloc((float**)&d_out, ARRAY_BYTES);
    cudaMalloc((float**)&d_in2, ARRAY_BYTES);
	cudaMalloc((float**)&d_out2, ARRAY_BYTES);
    cudaMalloc((float**)&d_in3, ARRAY_BYTES);
	cudaMalloc((float**)&d_out3, ARRAY_BYTES);
	cudaMemcpy(d_in, h_in, ARRAY_BYTES, cudaMemcpyHostToDevice);
    cudaMemcpy(d_in2, h_in, ARRAY_BYTES, cudaMemcpyHostToDevice);
    cudaMemcpy(d_in3, h_in, ARRAY_BYTES, cudaMemcpyHostToDevice);
    global_scan<<<1, ARRAY_SIZE>>>(d_out, d_in);
    sharemem_scan<<<1, ARRAY_SIZE, ARRAY_SIZE*sizeof(float)>>>(d_out2, d_in2);
    blelloch_scan<<<1, ARRAY_SIZE, ARRAY_SIZE*sizeof(float)>>>(d_out3, d_in3);
    cudaMemcpy(h_out, d_out, ARRAY_BYTES, cudaMemcpyDeviceToHost);
    cudaMemcpy(h_out2, d_out2, ARRAY_BYTES, cudaMemcpyDeviceToHost);
    cudaMemcpy(h_out3, d_out3, ARRAY_BYTES, cudaMemcpyDeviceToHost);

    printf("global scanf:\n");
    for(int i = 0; i < ARRAY_SIZE; i++){
        printf("%f", h_out[i]);
        printf(((i % 4)!= 3) ? "\t" : "\n");
    }
    printf("shared memory scanf:\n");
    for(int i = 0; i < ARRAY_SIZE; i++){
        printf("%f", h_out2[i]);
        printf(((i % 4) != 3) ? "\t" : "\n");
    }
    printf("blelloch scanf:\n");
    for(int i = 0; i < ARRAY_SIZE; i++){
        printf("%f", h_out3[i]);
        printf(((i % 4) != 3) ? "\t" : "\n");
    }

    cudaFree(d_in);
	cudaFree(d_out);
    cudaFree(d_in2);
	cudaFree(d_out2);
    cudaFree(d_in3);
	cudaFree(d_out3);
	cudaDeviceReset();

    return 0;
}